from llmtcg.agents.base import Agent
from llmtcg.chat.model import ChatContext, Role

from loguru import logger

from dataclasses import dataclass, field
from typing import List, Tuple, Optional
from tempfile import mkstemp, mkdtemp
from subprocess import run, PIPE, CompletedProcess, STDOUT, TimeoutExpired, CalledProcessError
from os import path, makedirs, environ, system
import re

from shutil import rmtree
import uuid


CCODER_SYSTEM_PROMPT = """
You are a professional programmer and are supposed to write a program to
achieve the given requirement. Your response should start with a brief
introduction of how the program works, and then the generated program.
The generated program should be surrounded by "```c" before it and
"```" after it.

You should always carefully check whether the header files are properly included
if you have used some function that are declared in them. For example, if you have used
function "strlen", "strcpy", "strcmp", "strcat", "memset", "memcpy", "memcmp", or "memmove",
you need to add "#include <string.h>" at the beginning.

Moreover, use fprintf(stderr) to get the values of relevant variables at important location in the code, 
such as before executing the assert function.
"""

CCODER_OPTEE_SYSTEM_PROMPT = """
You are a professional programmer and are supposed to write a OPTEE test case to achieve the given requirement. 
Your response should start with a brief introduction of how the program works, and then the generated program. 
The generated program should contain: 
    1. Client Application side C code which should be surrounded by "```ca" before it and "```" after it;
    2. The Client Application side code should contain a main function.
    3. Trusted Application side C code which should be surrounded by "```ta" before it and "```" after it;

You should always carefully check whether the header files are properly included if you have used some function that are declared in them. For example, if you have used the function "strlen", you need to add "#include <string.h>" at the beginning.

Moreover, print the values of relevant variables at important location in the code, such as before executing the assert function.
"""

DRAWAPI_SYSTEM_PROMPT="""
You are a professional programmer and are supposed to write a program to achieve the given requirement. Your response should start with a brief introduction of how the program works, and then the generated program.

The generated program should be surrounded by "```cpp" before it and "```" after it.
"""


CCODER_RETRY_PROMPT = """
Your generated program encounters the following problem: {problem_description},
detailed message is shown as follows.

```
{feedback}
```

Please fix the previous program and re-generate the correct one that matches the original
requirement.
"""

CCODER_FAKE_RESP = """
The program is generated as follows,

```c
{program}
```
"""


class CodingFailed(Exception):
    pass


@dataclass
class CCoderAgent(Agent):

    compiler: str = "gcc"
    compiler_args: List[str] = field(default_factory=list)
    max_retries: int = 3
    with_history: bool = True

    _last_context: Optional[ChatContext] = None

    @property
    def system_prompt(self) -> str:
        return CCODER_SYSTEM_PROMPT

    @property
    def last_context(self) -> ChatContext:
        assert self._last_context is not None, f"no last context found"
        return self._last_context

    def validate_program(self, program: str, tfname: str) -> Tuple[bool, bool, str]:
        """
        We use a inherited temp file name tfname from outside world, so that every time
        we compile/run the program, the produced error message has the same source file
        name.
        """
        with open(tfname, "w+") as tf:
            tf.write(program)

        outname = f"{tfname}.out"
        cproc: CompletedProcess = run(
            [self.compiler, *self.compiler_args, tfname, "-o", outname, "-lssl", "-lcrypto", "-lpthread", "-lrt", "-lm",],
            stdin=PIPE, stdout=PIPE, stderr=PIPE
        )
        if cproc.returncode != 0:
            return False, False, cproc.stderr.decode()

        temp_dir = mkdtemp(prefix="llmtcg-coder-")
        makedirs(temp_dir, exist_ok=True)

        try:
            # permission = environ.get("root_pwd")
            rproc = run(
                [outname],
                cwd=temp_dir, stdin=PIPE, capture_output=True, start_new_session=True, timeout=20
            )
            rmtree(temp_dir, ignore_errors=True)
            if rproc.returncode != 0:
                if rproc.returncode == -11:
                    msg = """ Running this program caused a segmentation fault. 
                    Please modify the code to prevent passing NULL pointers to the function. 
                    Also, avoid abnormal memory handling behaviors, such as dereferencing a NULL pointer, 
                    using wild pointers, releasing unallocated memory, etc. 
                    Avoid allocating too large local variables or using infinite recursion to prevent stack overflow. 
                    """
                elif rproc.returncode == -8:
                    msg = """Running this program results in SIGFPE.
                    Please modify the code to correct the floating-point number exception, 
                    such as avoiding division by zero operations.
                    """
                else:
                    msg = rproc.stdout.decode() + rproc.stderr.decode()
                return True, False, msg
        except TimeoutExpired:
            msg = """ Running this program lead to a timeout problem, probably because
            the program is waiting for something from another process or stdin that will
            never come.
            """
            rmtree(temp_dir, ignore_errors=True)
            return True, False, msg
        except UnicodeDecodeError:
            rmtree(temp_dir, ignore_errors=True)
            return True, False, "failed to decode the output of the test case"
        except CalledProcessError as e:
            if e.returncode == 13 or e.returncode == 126 or "Permission denied" in e.stderr:
                return True, True, "failed because of permission denied but we kept this case"

        return True, True, ""

    def get_program(self, requirement: str) -> str:
        if self.with_history:
            return self.__get_program_with_history(requirement)
        else:
            return self.__get_program_without_history(requirement)

    def __extract_program(self, response: str) -> str:
        try:
            start_pos = response.index("```c")
            end_pos = response.index('```', start_pos + 1)
        except ValueError:
            raise CodingFailed("cannot extract prgoram from response")

        program = response[start_pos + 4: end_pos].strip()
        return program

    def __get_program_with_history(self, requirement: str) -> str:
        _, tfname = mkstemp(prefix="llmtcg-agent-ccoder", suffix=".c")

        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        try:
            resp = chat_ctx.chat(f"Please write a c program according to the given requirement: {requirement}")
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}")

        nretry = 0
        last_error = ""

        while True:
            try:
                start_pos = resp.index("```c")
                end_pos = resp.index('```', start_pos + 1)
            except ValueError:
                logger.debug(resp)
                raise CodingFailed("cannot extract prgoram from response")

            program = resp[start_pos + 4: end_pos].strip()
            compilable, runnable, msg = self.validate_program(program, tfname)

            if compilable and runnable:
                """ We are good ;) """
                logger.info(f"c program generated with {nretry} retries")
                return program

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit")

            """ Update last error """
            last_error = msg

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                raise CodingFailed(f"retry exceed the limitation")

            # logger.debug('------ old program ------')
            # logger.debug(program)

            if not compilable:
                # logger.debug('----- compiling error -----')
                # logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program you wrote, it cannot be compiled.
                {self.compiler} reported the following error:

                    {msg.strip()}
                """
            else:
                # logger.debug('----- running error -----')
                # logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program you wrote, it can be successfully
                compiled, but running it failed. The contents of stdout and stderr are as follows:

                    {msg.strip()}
                """

            try:
                resp = chat_ctx.chat(new_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}")

    def __get_program_without_history(self, requirement: str) -> str:
        _, tfname = mkstemp(prefix="llmtcg-agent-ccoder", suffix=".c")

        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        initial_prompt = f"Please write a c program according to the given requirement: {requirement}"
        try:
            resp = chat_ctx.chat(initial_prompt)
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}")

        nretry = 0
        last_error = ""

        while True:
            program = self.__extract_program(resp)
            compilable, runnable, msg = self.validate_program(program, tfname)

            if compilable and runnable:
                """ We are good ;) """
                logger.info(f"c program generated with {nretry} retries")
                return program

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit")

            """ Update last error """
            last_error = msg

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                raise CodingFailed(f"retry exceed the limitation")

            if not compilable:
                problem_description = f"the program cannot be compiled, {self.compiler} reported the following error"
            else:
                problem_description = \
                    "the program can be successfully compiled, " + \
                    "but running it lead to the following error"

            retry_prompt = CCODER_RETRY_PROMPT.format(
                problem_description=problem_description,
                feedback=msg
            )

            try:
                chat_ctx.clear()
                chat_ctx.history.append((Role.USER, initial_prompt))
                chat_ctx.history.append((Role.ASSISTANT, CCODER_FAKE_RESP.format(program=program)))
                resp = chat_ctx.chat(retry_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}")


@dataclass
class CKernelCoderAgent(CCoderAgent):

    ''' for each api, we generate n scenes. For each scene, we try many rounds
        to make sure the generated program is compilable and runnable.

        I oserved that the compile erros meet in this scene are often meet again
        in the following scenes, so for the same api with different scenes, we
        may pass previous compiling errors to avoid the same errors in the future.
    '''

    def get_program(self, requirement: str) -> str:
        _, tfname = mkstemp(prefix="llmtcg-agent-ccoder", suffix=".c")

        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        try:
            resp = chat_ctx.chat(f"Please write a c program according to the given requirement: {requirement}")
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}", chat_ctx)

        nretry = 0
        last_error = ""

        while True:
            try:
                start_pos = resp.index("```c")
                end_pos = resp.index('```', start_pos + 1)
            except ValueError:
                logger.debug(resp)
                raise CodingFailed("cannot extract prgoram from response", chat_ctx)

            program = resp[start_pos + 4: end_pos].strip()
            compilable, runnable, msg = self.validate_program(program, tfname)

            if compilable and runnable:
                """ We are good ;) """
                logger.info(f"c program generated with {nretry} retries")
                return program

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit", chat_ctx)
            
            """ Update last error """
            last_error = msg

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                raise CodingFailed(f"retry exceed the limitation", chat_ctx)

            logger.debug('------ old program ------')
            logger.debug(program)

            if not compilable:
                logger.debug('----- compiling error -----')
                logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program your written, it cannot be compiled.
                {self.compiler} report the following error:

                    {msg.strip()}
                """
            else:
                logger.debug('----- running error -----')
                logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program your written, it can be successfully
                compiled, but running it leads to the following error:

                    {msg.strip()}
                """

            try:
                resp = chat_ctx.chat(new_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}", chat_ctx)

    def validate_program(self, program: str, tfname: str) -> Tuple[bool, bool, str]:
        temp_dir = mkdtemp(prefix="llmtcg-coder-")
        makedirs(temp_dir, exist_ok=True)

        tfname = tfname.replace("/tmp", temp_dir)
        print("file_name:" + tfname)

        with open(tfname, "w+") as tf:
            tf.write(program)

        # generate a Makefile for each generated c file
        outname = tfname.split('/')[-1]
        outname = outname.replace(".c", ".o")
        mk_file = path.join(temp_dir, "Makefile")
        content = f"""obj-m += {outname}\n\nall:\n\tmake -C ~/WSL2-Linux-Kernel M=`pwd` EXTRA_CFLAGS='-w' modules\nclean:\n\tmake -C ~/WSL2-Linux-Kernel M=`pwd` clean"""
        with open(mk_file, "w+") as mkf:
            mkf.write(content)

        try:
            # compile linux kernel module and generate ko file
            make_proc: CompletedProcess = \
            run(["make"], cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)

        except Exception as e:
            print("Exception when compling kernel module error:"+str(e))
        print("make return code: "+str(make_proc.returncode))

        if make_proc.returncode != 0:
            msg = make_proc.stderr.decode()
            # extract the first error message
            errors = re.findall(r'error:.*\n.*\n.*', msg)
            if errors:
                msg = errors[0]
            return False, False, msg

        # free resources
        def free_resources() -> None:
            try:
                run("sync", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
                run(f"sudo -S echo 3 > /proc/sys/vm/drop_caches",
                    input=permission+'\n', text=True, start_new_session=True,
                    shell=True, stdout=PIPE, stderr=PIPE)
                run("sleep 1", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
            except Exception as e:
                logger.error(f"failed to free resources: {e}")

        try:
            # load generated ko file by insmod
            ko_file = tfname.replace(".c", ".ko")
            permission = environ.get("root_pwd")
            insmod_proc: CompletedProcess = run(
                f"sudo -S insmod {ko_file}", input=permission+'\n', text=True,
                cwd=temp_dir, start_new_session=True, stdout=PIPE, stderr=PIPE,
                shell=True, timeout=20)
            print("insmod return code: "+str(insmod_proc.returncode))

            if insmod_proc.returncode != 0:
                free_resources()
                rmtree(temp_dir, ignore_errors=True)
                msg = f"insmod error: return code is {insmod_proc.returncode}: " + insmod_proc.stderr
                return True, False, msg

        # meet unexpected insmod timeout error, but disappear after reboot
        # So keep the generated c program and reboot once addressing several scenes
        except TimeoutExpired as ex:
            msg = "insmod error:"+str(ex)
            free_resources()
            return True, True, msg

        except Exception as e:
            print("insmod error:"+str(e))
            free_resources()
            rmtree(temp_dir, ignore_errors=True)
            return True, False, "insmod error: " + str(e)

        try:
            # unload ko file by rmmod
            rmmod_proc: CompletedProcess = run(
                f"sudo -S rmmod {ko_file}", input=permission+'\n', text=True,
                cwd=temp_dir, stdout=PIPE, stderr=PIPE,
                shell=True, start_new_session=True, timeout=5)
            print("rmmod return code: "+str(rmmod_proc.returncode))
            free_resources()

            if rmmod_proc.returncode != 0:
                rmtree(temp_dir, ignore_errors=True)
                msg = f"rmmod return code is {rmmod_proc.returncode}: " + rmmod_proc.stderr
                return True, False, msg

        except TimeoutExpired as ex:
            msg = "rmmod error:"+str(ex)
            free_resources()
            return True, True, msg

        except Exception as e:
            free_resources()
            rmtree(temp_dir, ignore_errors=True)
            return True, False, "rmmod error: " + str(e)

        return True, True, ""


@dataclass
class COPTEECoderAgent(Agent):

    compiler_args: List[str] = field(default_factory=list)
    max_retries: int = 5
    with_history: bool = True
    _last_context: Optional[ChatContext] = None

    @property
    def system_prompt(self) -> str:
        return CCODER_OPTEE_SYSTEM_PROMPT

    @property
    def last_context(self) -> ChatContext:
        assert self._last_context is not None, f"no last context found"
        return self._last_context
    
    def __extract_program(self, response: str):
        try:
            start_pos = response.index("```ca")
            end_pos = response.index("```", start_pos + 1)
            ca_program = response[start_pos + 5: end_pos].strip()
        except ValueError:
            raise CodingFailed("cannot extract prgoram from response")

        try:
            start_pos = response.index("```ta")
            end_pos = response.index("```", start_pos + 1)
            ta_program = response[start_pos + 5: end_pos].strip()
        except ValueError:
            raise CodingFailed("cannot extract prgoram from response")

        return ca_program, ta_program

    def validate_program(self, temp_dir: str, ca_program: str, ta_program: str, binary_name:str) -> Tuple[bool, str]:
        
        try:
            ca_propgram_path = path.join(temp_dir, "ca", "ca.c")
            ta_propgram_path = path.join(temp_dir, "ta", "ta.c")
            with open(ca_propgram_path, "w+") as fprog:
                fprog.write(ca_program)
            with open(ta_propgram_path, "w+") as fprog:
                fprog.write(ta_program)

            make_proc: CompletedProcess = \
                run(["make", "CROSS_COMPILE=aarch64-linux-gnu-", "-lpthread"], 
                    cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)
        
            print("make return code: "+str(make_proc.returncode))
            if make_proc.returncode != 0:
                msg = make_proc.stderr.decode()
                errors = re.findall(r'error:.*\n.*\n.*', msg)
                if errors:
                    msg = errors[0]
                run(["make", "clean"], cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)
                return False, make_proc.stderr.decode()
            
            binary_path = path.join(temp_dir, "ca", binary_name)
            if not path.exists(binary_path):
                print("Failed to generate binary file")
                run(["make", "clean"], cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)
                return False, "Failed to genrate binary file"
            
            run(["make", "clean"], cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)
            
        except Exception as e:
            print("Exception when compling, error:" + str(e))
            run(["make", "clean"], cwd=temp_dir, stdin=PIPE, stdout=PIPE, stderr=PIPE)
            return False, str(e)
                
        return True, ""
    

    def init_temp_files(self, temp_dir:str, _uuid: uuid.UUID, c_define: str, binary_name: str) -> None:
        def _init_ca_makefile() -> None:
            ca_makefile = path.join(temp_dir, "ca", "Makefile")
            with open(ca_makefile, 'r') as file:
                content = file.read()
            new_content = content.replace('BINARY = _to_be_filled', f"BINARY = {binary_name}")
            with open(ca_makefile, 'w') as file:
                file.write(new_content)
        
        def _init_ta_makefile() -> None:
            ta_makefile = path.join(temp_dir, "ta", "Makefile")
            with open(ta_makefile, 'r') as file:
                content = file.read()
            new_content = content.replace('BINARY = _to_be_filled', f"BINARY = {_uuid}")
            with open(ta_makefile, 'w') as file:
                file.write(new_content)
        
        def _init_ta_header() -> None:
            ta_header = path.join(temp_dir, "ta", "user_ta_header_defines.h")
            with open(ta_header, 'r') as file:
                content = file.read()
            new_content = content.replace('#define TA_UUID _to_be_filled', f"{c_define}")
            with open(ta_header, 'w') as file:
                file.write(new_content)

        def _init_cmakelist() -> None:
            cmakelist = path.join(temp_dir, "CMakeLists.txt")
            with open(cmakelist, 'r') as file:
                content = file.read()
            new_content = content.replace('project (_to_be_filled C)', f"project ({binary_name} C)")
            with open(cmakelist, 'w') as file:
                file.write(new_content)
        
        _init_cmakelist()
        _init_ca_makefile()
        _init_ta_header()
        _init_ta_makefile()
    

    def get_program(self, requirement: str, _uuid: uuid.UUID, c_define: str, binary_name: str):
        temp_dir = mkdtemp(prefix="llmtcg-optee-coder-")
        makedirs(temp_dir, exist_ok=True)
        system(f"cp -r testcase_template/* {temp_dir}")
        self.init_temp_files(temp_dir, _uuid, c_define, binary_name)
        
        chat_ctx = ChatContext(self.chat_model)
        chat_ctx.task_identifier = "CCoderAgent::getProgram"
        chat_ctx.system_prompt = self.system_prompt

        self._last_context = chat_ctx

        try:
            resp = chat_ctx.chat(requirement)
        except Exception as ex:
            raise CodingFailed(f"chat API failed because: {ex}")

        nretry = 0
        last_error = ""

        while True:
            ca_program, ta_program = self.__extract_program(resp)
            compilable, msg = self.validate_program(temp_dir, ca_program, ta_program, binary_name)

            if compilable:
                """ We are good ;) """
                logger.info(f"testing program generated with {nretry} retries")
                return temp_dir

            if msg == last_error:
                raise CodingFailed(f"the exact error message reported again, quit")

            """ Update last error """
            last_error = msg

            """ Make a retry ;) """
            nretry += 1

            if nretry > self.max_retries:
                system(f"rm -rf {temp_dir}")
                raise CodingFailed(f"retry exceed the limitation, remove temp directory")

            logger.debug('------ old program ------')
            logger.debug('```ca```')
            logger.debug(ca_program)
            logger.debug('```ta```')
            logger.debug(ta_program)

            if not compilable:
                logger.debug('----- compiling error -----')
                logger.debug(msg.strip())
                new_prompt = f"""
                Please correct the previous program you wrote, it cannot be compiled.
                compiling reported the following error:

                    {msg.strip()}
                """

            try:
                resp = chat_ctx.chat(new_prompt)
            except Exception as ex:
                raise CodingFailed(f"chat API failed because: {ex}")
